import matplotlib.pyplot as plt
import numpy as np
import matplotlib
import os
import matplotlib.ticker as mtick
matplotlib.use('TkAgg')
matplotlib.rc("font", family='simsun,', weight="bold", size=12)
data_name = ["total", "hole", "crack", "normal"]
metric_name = ["loss", "acc", "precision", "recall", "f1"]


def get_info(path):
    """
    paramter: path is the path of txt file which records the result information of validation dataset
    during the training process
    return: return the result of all model on validation dataset
    """
    all_model_data = {}
    for package in os.listdir(path):
        if package.endswith('.jpg') or package.endswith('png') or package.endswith(".docx") or package.startswith(
                "result"):
            continue
        one_model_data = {}
        val_info_path = os.path.join(path, package) + '/' + "val_info.txt"
        val_info = open(val_info_path, 'r')
        val_info_lines = val_info.readlines()
        val_total_dict = {"loss": [], "acc": [], "precision": [], "recall": [], "f1": []}
        val_hole_dict = {"loss": [], "acc": [], "precision": [], "recall": [], "f1": []}
        val_crack_dict = {"loss": [], "acc": [], "precision": [], "recall": [], "f1": []}
        val_normal_dict = {"loss": [], "acc": [], "precision": [], "recall": [], "f1": []}
        dict_name = [val_total_dict, val_hole_dict, val_crack_dict, val_normal_dict]
        for i in range(4):
            for idx, j in enumerate(range(i * 5, i * 5 + 5)):
                dict_name[i][metric_name[idx % 5]] = val_info_lines[j].strip().split(":")[1]
        one_model_data['total'] = val_total_dict
        one_model_data["hole"] = val_hole_dict
        one_model_data["crack"] = val_crack_dict
        one_model_data["normal"] = val_normal_dict
        all_model_data[package] = one_model_data
    return all_model_data


def show_curve_zh(val_dict, title, ax, symbol):
    """
    :param
    """
    color = ["blue", "black", "red", "purple", "green"]  # the color of curve
    legend = []  # a list to record the legend of every curve
    x = list(range(1, 101))
    for idx, (key, value) in enumerate(val_dict.items()):
        value_list = [float(v) for v in value[1:-1].split(",")]
        ax.plot(x, value_list, linewidth=2, color=color[idx])
        legend.append(key)
    # ax.axhline(y=0.95, linewidth=2, linestyle="--")
    # ax.annotate(s="0.95", xy=(0, 0.92), xytext=(0, 10), textcoords='offset points', fontsize=30, color="red")
    ax.annotate(s=f"{symbol}", xy=(95, 0.5), xytext=(0, 4), textcoords='offset points', fontsize=40, color="black")
    ax.legend(legend, fontsize=26)
    ax.set_xlabel("训练轮数/epoch", fontsize=40)
    ax.set_ylabel(title, fontsize=40)
    ax.set_yticks(np.arange(0, 1.1, 0.1))
    ax.set_xticks(np.arange(0, 110, 10))
    ax.set_xticklabels(np.arange(0, 110, 10), fontsize=35)
    ax.set_yticklabels(np.arange(0, 1.1, 0.1), fontsize=35)
    ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))


def show_curve(val_dict, title, path):
    color = ["blue", "black", "red", "purple", "green"]
    legend = []
    figure=plt.figure(figsize=(8,6),dpi=1500)
    x = list(range(1, 101))
    for idx, (key, value) in enumerate(val_dict.items()):
        value_list = [float(v) for v in value[1:-1].split(",")]
        plt.plot(x, value_list, linewidth=1, color=color[idx])
        legend.append(key)
    plt.legend(legend, fontsize=20)
    plt.xlabel("训练轮数/epoch", fontsize=20)
    plt.ylabel(title, fontsize=20)
    plt.yticks(np.arange(0, 1.05, 0.1), fontsize=20)
    plt.xticks(np.arange(0, 110, 10), fontsize=20)
    plt.axhline(y=1, linewidth=1, linestyle="--")
    # plt.savefig(f"{path}/result_curve/{title}.png")
    plt.show()


def compute_mean(data):
    result = []
    for key, value in data.items():
        value = [float(v) for v in value[1:-1].split(",")]
        np_value = np.array(value)
        np_value = np_value[80:101]
        np_value.sort()
        np_value = np_value[1:-1]
        mean_value = np_value.mean()
        std_value = np_value.std()
        result.append([mean_value, std_value])
    return tuple(result)


if __name__ == "__main__":
    path = "2700+/train_data+LR"
    all_model_data = get_info(path)
    loss_dict = {}
    acc_dict = {}
    hole_dict = {}
    crack_dict = {}
    normal_dict = {}
    total_dict = {}
    ###############################若需要绘制其指标的图片，只需要改变value["xxxx"]中的内容即可：loss,precision,acc,precision,recall,f1
    element = "f1"
    for model_name, model_data in all_model_data.items():
        print(f"####################{model_name}###################")
        for defect_name, value in model_data.items():
            loss_dict[model_name] = value["loss"]
            acc_dict[model_name] = value["acc"]
            if defect_name == "total":
                total_dict[model_name] = value[element]
            elif defect_name == "hole":
                hole_dict[model_name] = value[element]
            elif defect_name == "crack":
                crack_dict[model_name] = value[element]
            elif defect_name == "normal":
                normal_dict[model_name] = value[element]
            loss, acc, precision, recall, f1 = compute_mean(value)
            print(f"{defect_name} ||  acc:[{acc[0]:.04f}, {acc[1]:.04f}] f1:[{f1[0]:.04f},{f1[1]:.04f} ]")
    symbols = ["(a)", "(b)", "(c)", "(d)"]
    Dict_data = [hole_dict, crack_dict, normal_dict, total_dict]
    ylabels = ["气孔图片分类 F1-score", "裂纹图片分类 F1-score", "正常图片分类 F1-score", "F1-score 宏平均值"]
    # fig, axes = plt.subplots(2, 2, figsize=(24, 18), dpi=1000)
    # for i in range(0, 4):
    #     show_curve_zh(Dict_data[i], ylabels[i], axes[int(i / 2), int(i % 2)], symbols[i])
    # plt.savefig(f"{path}\\result_curve\\'zh'.png")

    show_curve(acc_dict, "分类准确度accuracy", path)
